package com.shanzhi.dify.utils;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.TypeReference;
import com.shanzhi.dify.entity.ChunkEntity;
import com.shanzhi.dify.entity.DocumentFileEntity;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 用于Dify的静态解析数据的工具类
 */
public class ParseDataUtils {


    /**
     * dify数据解析-聊天记录-ID
     */
    public static String parseDataForGetLastConversationId(String jsonData) {
        // 解析JSON数据
        JSONObject jsonObject = JSONObject.parseObject(jsonData);
        JSONArray dataArray = jsonObject.getJSONArray("data");

        if (dataArray != null && !dataArray.isEmpty()) {
            // 获取最后一个对话项
            JSONObject lastConversation = dataArray.getJSONObject(0);
            return lastConversation.getString("id");
        }
        return null;
    }


    /**
     * 从知识库搜索结果中提取所有segment的content内容
     *
     * @param jsonString 知识库搜索的完整JSON响应
     * @return 所有记录的segment content内容列表
     */
    public static List<String> extractSegmentContent(String jsonString) {
        List<String> result = new ArrayList<>();
        JSONObject jsonObject = JSONObject.parseObject(jsonString);
        JSONArray records = jsonObject.getJSONArray("records");

        if (records != null && !records.isEmpty()) {
            for (int i = 0; i < records.size(); i++) {
                JSONObject record = records.getJSONObject(i);
                JSONObject segment = record.getJSONObject("segment");
                String content = segment.getString("content");
                if (content != null) {
                    result.add(content);
                }
            }
        }

        return result;
    }


    /**
     * 从知识库搜索结果中提取所有segment的content内容，并按文档分组封装为DocumentFileEntity对象
     *
     * @param jsonString 知识库搜索的完整JSON响应
     * @return 所有记录按文档分组的DocumentFileEntity列表
     */
    public static List<DocumentFileEntity> extractSegmentContentForDocument(String jsonString) {
        Map<String, DocumentFileEntity> documentMap = new HashMap<>();
        JSONObject jsonObject = JSONObject.parseObject(jsonString);
        JSONArray records = jsonObject.getJSONArray("records");

        if (records != null && !records.isEmpty()) {
            for (int i = 0; i < records.size(); i++) {
                JSONObject record = records.getJSONObject(i);
                JSONObject segment = record.getJSONObject("segment");

                String content = segment.getString("sign_content");
                String segmentId = segment.getString("id"); // chunkId
                String documentId = segment.getString("document_id");

                JSONObject documentObj = segment.getJSONObject("document");
                String fileName = documentObj != null ? documentObj.getString("name") : null;

                if (segmentId != null && documentId != null && fileName != null) {
                    // 如果文档不存在则创建新文档实体
                    if (!documentMap.containsKey(documentId)) {
                        DocumentFileEntity documentEntity = new DocumentFileEntity();
                        documentEntity.setDocumentId(documentId);
                        documentEntity.setFileName(fileName);
                        documentEntity.setChunkEntityList(new ArrayList<>());
                        documentMap.put(documentId, documentEntity);
                    }

                    // 创建ChunkEntity并添加到对应文档的chunkEntityList中
                    ChunkEntity chunkEntity = new ChunkEntity();
                    chunkEntity.setChunkId(segmentId);
                    chunkEntity.setContent(content);
                    chunkEntity.setSimilarScore(record.getDoubleValue("score"));

                    documentMap.get(documentId).getChunkEntityList().add(chunkEntity);
                }
            }
        }

        return new ArrayList<>(documentMap.values());
    }


    /**
     * 获取知识库ID
     * @param datasetJson
     * @return
     */
    public static String getDatasetId(String datasetJson) {
        Map<String, Object> stringObjectMap = ParseDataUtils.parseJson(datasetJson, new TypeReference<Map<String, Object>>() {
        });
        return stringObjectMap.get("id").toString();

    }


    /**
     * 通用JSON解析方法
     *
     * @param json  JSON字符串
     * @param clazz 目标类型Class
     * @param <T>   泛型类型
     * @return 解析后的对象
     */
    public static <T> T parseJson(String json, Class<T> clazz) {
        return JSON.parseObject(json, clazz);
    }

    /**
     * 通用JSON解析方法（支持复杂泛型类型）
     *
     * @param json          JSON字符串
     * @param typeReference 类型引用
     * @param <T>           泛型类型
     * @return 解析后的对象
     */
    public static <T> T parseJson(String json, TypeReference<T> typeReference) {
        return JSON.parseObject(json, typeReference);
    }


}
